feat(pipeline): Add inference and lineage step types - #6224
Conversation
| return _SHAPE_CACHE[cache_key] | ||
|
|
||
|
|
||
| def validate_step_arguments( |
There was a problem hiding this comment.
why do we need all of this additional validation here? Other steps do not have this explicit validation. How are these steps different from other steps?
| def __init__( | ||
| self, | ||
| name: str, | ||
| arguments: Dict[str, Any], |
There was a problem hiding this comment.
This is not the right implementation for any of these steps. It will be very difficult to construct these arguments manually. We need to use the existing pysdk constructs and pass them as arguments. Please see how Training/Model steps are implemented and follow that pattern here. You must use step_args from PipelineSession instead of raw arguments: dict
SDK primitives exists for all four steps, and it eliminates the entire _argument_validation.py machinery
| return get_execution_role() | ||
|
|
||
|
|
||
| def test_lineage_step_execute_end_to_end(sagemaker_session, pipeline_session, role): |
There was a problem hiding this comment.
please add integ tests for other steps as well
Adds four pipeline step classes -- EndpointConfigStep, EndpointStep, InferenceComponentStep and LineageStep -- along with their StepTypeEnum values and package exports. Each step takes step_args captured under a PipelineSession, following the convention used by TrainingStep and ModelStep: - Session.endpoint_from_production_variants, Session.create_endpoint and Session.create_inference_component route their service calls through _intercept_create_request. Under a plain Session this is a pass-through; under a PipelineSession the request is captured and no service call is made. - LineageStep records one lineage entity per step, with step_args from Action.create(), Artifact.create(), Context.create() or Association.create(). Record._invoke_api captures those four calls under a PipelineSession. An association can reference an entity created by another step via Steps.<name>.ActionArns['<entity>']. - Each step validates the provenance of its step_args with validate_step_args_input, rejecting a wrong producer or a raw dict. Note for chaining: the resource created by an EndpointConfig step carries an execution-unique name suffix, so a downstream step must reference it as config_step.properties.EndpointConfigName rather than the requested name. Tested with 29 unit tests plus two integration tests that exercise all four step types end to end against the service. --- X-AI-Prompt: Add the new inference and lineage pipeline step types to the public SDK, using step_args captured via PipelineSession per review feedback, and verify them with integration tests X-AI-Tool: kiro-cli
00feb86 to
00df5b4
Compare
Add 4 pipeline step classes:
Design: each step takes step_args captured under a PipelineSession, following the convention used by TrainingStep and ModelStep. Session.endpoint_from_production_variants, Session.create_endpoint and Session.create_inference_component route through _intercept_create_request -- a pass-through under a plain Session, capturing the request without a service call under a PipelineSession. LineageStep records one entity per step, with step_args from Action/Artifact/Context/Association.create(); an association can reference an entity created by another step via Steps..ActionArns['']. Each step validates step_args provenance with validate_step_args_input, rejecting a wrong producer or a raw dict.
Chaining note: the resource created by an EndpointConfig step carries an execution-unique name suffix, so a downstream step must reference it as config_step.properties.EndpointConfigName rather than the requested name. The Endpoint step does not suffix.
Retryability: only EndpointConfigStep is retryable. Cacheability: EndpointConfigStep and EndpointStep are structurally cacheable via cache_config.
Includes 29 unit tests and two integration tests covering all four step types end to end against the service (2 passed in 473.64s).
Supersedes #6140 -- same feature, reworked to step_args per review feedback.
Issue #, if available:
Description of changes:
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.